import numpy as np
from PIL import Image, ImageDraw
import os

PATH = os.path.join('dataSet', 'dbrhd', 'extend_3.tra')  # the path of the data set file
PREFIX = 'extend'  # prefix of the file name of the output image

NUM = (18, 9)  # columns * rows, or x * y
PADDING = 20
CELL_SIZE = (100, 100)  # pixels, size of one cell
IMG_SIZE = (NUM[0] * (CELL_SIZE[0] + 2*PADDING), NUM[1] * (CELL_SIZE[1] + 2*PADDING))  # 1px for cell border


def count_plus(count):
    """
    Plus 1 to the count, if count is more than NUM, the count will be reset to zero.
    The images will fill the columns first. If this column is full, then it turn to
    the next column.
    Return -> new_count, if_next_column, if_reset
    """
    if_next_column = False
    if_reset = False
    if count[1] < NUM[1]:
        count[1] += 1
    elif count[0] < NUM[0]:
        count[1] = 0
        count[0] += 1
        if_next_column = True
    else:
        count = [0, 0]
        if_reset = True

    return count, if_next_column, if_reset


def move_pos(points, count):
    """Move the points to their cell according to the count. Fill the column first,
    if a column is full, then jump to the next column. Expecting the param points are
    ndarray of 2 dimension."""
    if count[1] <= NUM[1]:
        bias_x = (CELL_SIZE[0] + 2*PADDING) * count[0] + PADDING
        bias_y = (CELL_SIZE[1] + 2*PADDING) * count[1] + PADDING
    else:
        raise ValueError('The count of row is too much. count:%s ' % str(count))

    points[:, 0] += bias_x
    points[:, 1] += bias_y

    # print('bias: ', (bias_x, bias_y), 'points: ', points)
    return points


def generate_filename(digit, prefix=PREFIX):
    """Generate a file name that does not exist."""
    n = 1
    while True:
        filename = r'%s%i-%i.png'%(prefix, digit, n)
        filename = os.path.join(os.getcwd(), 'dbrhd_images', filename)
        if os.path.isfile(filename):
            n += 1
        else:
            return filename


counts = [[0, 0] for i in range(10)]  # column * row, or x * y, count how many cells is used

raw_img = Image.new(size=IMG_SIZE, mode='RGB')
raw_dw = ImageDraw.Draw(raw_img)
# draw cells
w, h = CELL_SIZE[0] + 2*PADDING, CELL_SIZE[1] + 2*PADDING
for x in range(NUM[0]):
    raw_dw.line((x*w, 0, x*w, IMG_SIZE[1]), fill=(120,)*3, width=1)
for y in range(NUM[1]):
    raw_dw.line((0, y*h, IMG_SIZE[0], y*h), fill=(120,)*3, width=1)

imgs = [raw_img.copy() for i in range(10)]  # Images
dws = [ImageDraw.Draw(imgs[i]) for i in range(10)]  # Draws


with open(PATH) as f:
    lines = f.readlines()
    f.close()

print('generating pictures...')
for line in lines:
    data = line.split(',')
    digit = int(data.pop(-1))
    points = np.asarray(data, dtype=int).reshape([8, 2])  # convert <str> into <int>

    points = move_pos(points, counts[digit])

    # link sample points into trajectory
    dws[digit].line(tuple(points.reshape([16])), fill=(255, 255, 0), width=1)

    # draw sample points
    width = 4  # width of circle
    dws[digit].ellipse((tuple(points[0]-width), tuple(points[0]+width)), fill=(255,)*3)
    for i in range(1, len(points)):
        dws[digit].ellipse((tuple(points[i]-width), tuple(points[i]+width)), fill=(255, 0, 0))

    # update counts
    counts[digit], if_next_col, if_reset = count_plus(counts[digit])
    if if_reset:
        f_name = generate_filename(digit)
        img = imgs[digit].transpose(Image.FLIP_TOP_BOTTOM)
        img.save(f_name)

        imgs[digit] = raw_img.copy()
        dws[digit] = ImageDraw.Draw(imgs[digit])
    # print("counts: ", counts)


for i in range(10):
    if counts[i][0] != 0 or counts[i][1] != 0:
        f_name = generate_filename(i)
        img = imgs[i].transpose(Image.FLIP_TOP_BOTTOM)
        img.save(f_name)

        print(f_name)

print('All done!')



